// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef PINS_TESTS_P4_FUZZER_FUZZER_TESTS_H_
#define PINS_TESTS_P4_FUZZER_FUZZER_TESTS_H_

#include <functional>
#include <optional>
#include <string>
#include <tuple>
#include <vector>

#include "absl/container/btree_set.h"
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "gtest/gtest.h"
#include "p4/config/v1/p4info.pb.h"
#include "p4_fuzzer/fuzzer.pb.h"
#include "p4_fuzzer/fuzzer_config.h"
#include "p4_fuzzer/switch_state.h"
#include "p4_pdpi/ir.pb.h"
#include "thinkit/mirror_testbed_fixture.h"

namespace p4_fuzzer {

// Used for testing a specific milestone, ignoring MaskKnownFailures(), rather
// than everything, respecting MaskKnownFailures().
enum class Milestone {
  // Tests that the switch adheres to the minimum guarantees on resources
  // currently defined in
  kResourceLimits,
  kInvalidInputs,
};

inline std::vector<std::tuple<Milestone, std::string>>
MilestoneToStringMapping() {
  return {
      {Milestone::kResourceLimits, "resource_limits"},
      {Milestone::kInvalidInputs, "invalid_inputs"},
  };
}

struct FuzzerTestFixtureParams {
  thinkit::MirrorTestbedInterface *mirror_testbed;
  // The test assumes that the switch is pre-configured if no `gnmi_config` is
  // given (default), or otherwise pushes the given config before starting.
  std::optional<std::string> gnmi_config;
  p4::config::v1::P4Info p4info;
  // Maximum batch size of a single write request generated by the fuzzer.
  std::optional<int> max_batch_size;
  // The probability of performing a mutation on a given table entry.
  float mutate_update_probability = 0.1;
  // TODO: b/286413264 - Remove once multicast resources are modeled in p4.
  // The size of the built-in multicast group table.
  int multicast_group_table_size = 512;
  // Determines which type of issues the fuzzer detects. If left out, the fuzzer
  // will test everything, respecting MaskKnownFailures(). See declaration of
  // Milestone for more details.
  std::optional<Milestone> milestone;
  std::optional<std::string> test_case_id;
  // By default, the fuzzer attempts to exceed the listed resource guarantees on
  // all tables, allowing the switch to reject entries beyond those guarantees
  // with a RESOURCE_EXHAUSTED error.
  // This variable lets users specify a set of tables for which the fuzzer
  // should treat their resource guarantees as hard limits rather than trying to
  // go above them. If there are limitations or bugs on the switch causing it to
  // behave incorrectly when the resource guarantees of particular tables are
  // exceeded, this list can be used to allow the fuzzer to produce interesting
  // results in spite of this shortcoming.
  absl::btree_set<std::string>
      tables_for_which_to_not_exceed_resource_guarantees;
  // Fully qualified names of tables, actions, or match fields to skip during
  // fuzzing. Match field names should be prepended with the relevant table name
  // to uniquely identify them.
  // Users should ensure that any set that makes it impossible to generate a
  // valid table entry for a particular table contains the table itself.
  // TODO: Check the property above instead.
  absl::flat_hash_set<std::string> disabled_fully_qualified_names;
  // TODO: Fully qualified names of tables that do not support
  // MODIFY updates. This behaviour is not compliant with p4 runtime spec.
  absl::flat_hash_set<std::string> non_modifiable_tables;
  // The P4RT role the fuzzer should use.
  std::string p4rt_role;
  // A function for masking inequalities (due to known bugs) between entries
  // with the same TableEntryKey on the switch and in the fuzzer.
  std::optional<std::function<bool(const pdpi::IrTableEntry &,
                                   const pdpi::IrTableEntry &)>>
      TreatAsEqualDuringReadDueToKnownBug;
  // Ignores the constraints on tables listed when fuzzing entries.
  absl::flat_hash_set<std::string> ignore_constraints_on_tables;
  // This parameter makes the Fuzzer test NOT check and enforce fail-on-first
  // ordering in switch responses. By default, we expect such ordering.
  // TODO: Get this information from the P4Info instead of using a
  // parameter.
  bool do_not_enforce_fail_on_first_switch_ordering;
  // A function for masking any updates that should not be sent to the switch.
  std::function<absl::StatusOr<bool>(const FuzzerConfig&,
                                     const AnnotatedUpdate&)>
      IsBuggyUpdateThatShouldBeSkipped =
          [](const FuzzerConfig&, const AnnotatedUpdate& update) {
            return false;
          };

  // A function for modifying any `TableEntry` produced by the Fuzzer. Note that
  // mutations can override modified values.
  std::function<absl::Status(const pdpi::IrP4Info&, const SwitchState&,
                             p4::v1::TableEntry&)>
      ModifyFuzzedTableEntry =
          [](const pdpi::IrP4Info&, const SwitchState&, p4::v1::TableEntry&) {
            // By default, do nothing.
            return absl::OkStatus();
          };
  // A function for modifying any `MulticastGroupEntry` produced by the Fuzzer.
  // Note that mutations can override modified values.
  std::function<absl::Status(const pdpi::IrP4Info&, const SwitchState&,
                             p4::v1::MulticastGroupEntry&)>
      ModifyFuzzedMulticastGroupEntry = [](const pdpi::IrP4Info&,
                                           const SwitchState&,
                                           p4::v1::MulticastGroupEntry&) {
        // By default, do nothing.
        return absl::OkStatus();
      };
};

class FuzzerTestFixture
    : public testing::TestWithParam<FuzzerTestFixtureParams> {
protected:
  // Sets up the mirror test bed, then sets the test_case_id.
  void SetUp() override;

  // Resets switch state after a fatal failure by attempting to clear the switch
  // tables normally, falling back to rebooting the switch. Also runs the
  // standard mirror test bed tear down procedure.
  void TearDown() override;

  ~FuzzerTestFixture() override { delete GetParam().mirror_testbed; }
};

bool AbslParseFlag(absl::string_view milestone_text, Milestone *milestone,
                   std::string *error);
std::string AbslUnparseFlag(Milestone milestone);

} // namespace p4_fuzzer

#endif // PINS_TESTS_P4_FUZZER_FUZZER_TESTS_H_
